package lua;

import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
import redis.clients.jedis.*;
import redis.clients.jedis.exceptions.JedisClusterException;

import java.util.HashSet;
import java.util.Set;

/**
 * Java 调用限流脚本
 * 三、命令参数说明
 * 参数位置	含义	示例值	说明
 * KEYS[1]	限流键名	"ip:192.168.1.100"	格式建议："ip:实际IP地址"，确保唯一
 * ARGV[1]	时间窗口（秒）	60	表示 60 秒内进行限流
 * ARGV[2]	最大访问次数	10	时间窗口内允许的最大访问次数
 * ARGV[3]（可选）	每次访问的计数步长	1（默认）	特殊场景下可设为更大值（如一次请求算 5 次访问
 *
 * @author: MI
 * @email: 448341911@qq.com
 * @createTime: 2025/8/27 18:23
 * @updateUser: MI
 * @updateTime: 2025/8/27 18:23
 * @updateRemark: 修改内容
 * @version: 1.0
 */
public class ClusterIpRateLimiter {

    // 初始化 Lua 脚本（修正后的限流脚本）
    private static final String RATE_LIMIT_SCRIPT = "local expire = tonumber(ARGV[1])\n" +
            "local limit = tonumber(ARGV[2])\n" +
            "local step = tonumber(ARGV[3]) or 1\n" +
            "if #KEYS ~= 1 then\n" +
            "    return redis.error_reply(\"IP rate limiting requires exactly 1 key parameter\")\n" +
            "end\n" +
            "if #ARGV < 2 then\n" +
            "    return redis.error_reply(\"IP rate limiting requires at least 2 arguments\")\n" +
            "end\n" +
            "if not expire or expire <= 0 then\n" +
            "    return redis.error_reply(\"Expiration time must be a positive number\")\n" +
            "end\n" +
            "if not limit or limit <= 0 then\n" +
            "    return redis.error_reply(\"Maximum access count must be a positive number\")\n" +
            "end\n" +
            "if not step or step <= 0 then\n" +
            "    return redis.error_reply(\"Increment step must be a positive number\")\n" +
            "end\n" +
            "local num = tonumber(redis.call('INCRBY', KEYS[1], step))\n" +
            "if num == step then\n" +
            "    redis.call('EXPIRE', KEYS[1], expire)\n" +
            "end\n" +
            "if num > limit then\n" +
            "    return \"denied\"\n" +
            "else\n" +
            "    return \"allowed\"\n" +
            "end";

    private final JedisCluster jedisCluster;

    /**
     * 构造方法：初始化 Redis 集群连接（支持密码认证）
     *
     * @param clusterNodes      集群节点列表（格式：ip:port）
     * @param password          集群密码（可为 null，若集群无密码）
     * @param connectionTimeout 连接超时时间（毫秒）
     * @param soTimeout         读写超时时间（毫秒）
     * @param maxAttempts       集群最大重定向次数
     */
    public ClusterIpRateLimiter(Set<String> clusterNodes, String password,
                                int connectionTimeout, int soTimeout, int maxAttempts) {
        // 1. 解析集群节点
        Set<HostAndPort> nodes = new HashSet<>();
        for (String node : clusterNodes) {
            String[] ipPort = node.split(":");
            nodes.add(new HostAndPort(ipPort[0], Integer.parseInt(ipPort[1])));
        }

        GenericObjectPoolConfig<Connection> poolConfig = new GenericObjectPoolConfig<>();
        poolConfig.setMaxTotal(100);
        poolConfig.setMaxIdle(20);
        poolConfig.setMinIdle(5);
        poolConfig.setMaxWaitMillis(3000);


        // 2. 配置客户端参数（包含密码认证）
        JedisClientConfig clientConfig = DefaultJedisClientConfig.builder()
                .password(password)                         // 设置Redis密码
                .connectionTimeoutMillis(connectionTimeout) // 连接超时时间（毫秒）
                .socketTimeoutMillis(soTimeout)             // 读写超时时间（毫秒）
                .build();

        // 3. 初始化 JedisCluster（支持密码认证）
        this.jedisCluster = new JedisCluster(
                nodes,
                clientConfig,
                poolConfig
        );
    }

    /**
     * 检查 IP 是否允许访问（执行限流逻辑）
     *
     * @param ip            客户端 IP 地址
     * @param expireSeconds 限流时间窗口（秒）
     * @param maxAccess     时间窗口内最大访问次数
     * @param step          每次访问的计数步长（默认 1）
     * @return true=允许访问，false=拒绝访问
     */
    public boolean isAllowed(String ip, int expireSeconds, int maxAccess, Integer step) {
        // 1. 生成限流 Key（格式：ip:xxx.xxx.xxx.xxx）
        // 注意：集群环境下键会按哈希分配到不同节点，无需手动指定节点
        String limitKey = "ip:" + ip;

        // 2. 准备脚本参数
        String[] argv = new String[3];
        argv[0] = String.valueOf(expireSeconds); // ARGV[1]：过期时间
        argv[1] = String.valueOf(maxAccess);     // ARGV[2]：最大访问次数
        argv[2] = step != null ? String.valueOf(step) : "1"; // ARGV[3]：步长（默认1）

        try {
            // 3. 执行 Lua 脚本（集群环境下 JedisCluster 直接支持 eval 命令）
            Object result = jedisCluster.eval(
                    RATE_LIMIT_SCRIPT,  // Lua 脚本
                    1,                  // 键数量
                    limitKey,           // KEYS[1]
                    argv[0], argv[1], argv[2]  // ARGV 参数
            );

            // 4. 解析结果
            if ("allowed".equals(result)) {
                return true;
            } else if ("denied".equals(result)) {
                return false;
            } else {
                // 脚本返回错误信息（如参数错误）
                throw new RuntimeException("Rate limit error: " + result);
            }
        } catch (JedisClusterException e) {
            // 处理集群相关异常（如节点不可用、重定向失败等）
            throw new RuntimeException("Redis cluster error: " + e.getMessage(), e);
        }
    }

    /**
     * 关闭集群连接（应用退出时调用）
     */
    public void close() {
        if (jedisCluster != null) {
            jedisCluster.close();
        }
    }

    // 使用示例
    public static void main(String[] args) {
        // 1. 配置集群节点和密码
        Set<String> clusterNodes = new HashSet<>();
        clusterNodes.add("192.168.110.65:6379");
        clusterNodes.add("192.168.110.66:6379");
        clusterNodes.add("192.168.110.67:6379");
        String password = "seata@user123"; // 集群密码（替换为实际密码）

        // 2. 初始化限流工具
        ClusterIpRateLimiter limiter = new ClusterIpRateLimiter(
                clusterNodes,
                password,
                2000,  // 连接超时（毫秒）
                2000,  // 读写超时（毫秒）
                3      // 最大重定向次数
        );

        // 3. 测试 IP 限流：192.168.1.100 在 60 秒内最多访问 10 次
        String ip = "192.168.1.100";
        for (int i = 0; i < 12; i++) {
            boolean allowed = limiter.isAllowed(ip, 60, 10, 1);
            System.out.println("第 " + (i + 1) + " 次访问：" + (allowed ? "允许" : "拒绝"));
        }

        // 4. 关闭连接（实际应用中在销毁时调用）
        limiter.close();
    }
}
